// WebSocket API implementation for real-time updates
use axum::{
    extract::ws::{Message, WebSocket, WebSocketUpgrade},
    response::Response,
};
use codegraph_core::{EdgeId, NodeId, PerformanceMonitor};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::{broadcast, RwLock};
use tracing::{debug, error, info, warn};
use uuid::Uuid;

/// WebSocket connection manager
pub struct WebSocketManager {
    /// Active connections
    connections: Arc<RwLock<HashMap<Uuid, ConnectionInfo>>>,
    /// Broadcast channel for real-time updates
    update_sender: broadcast::Sender<GraphUpdateMessage>,
    /// Performance monitor
    performance_monitor: Arc<PerformanceMonitor>,
}

/// Information about a WebSocket connection
#[derive(Debug, Clone)]
struct ConnectionInfo {
    id: Uuid,
    user_id: Option<String>,
    subscriptions: Vec<SubscriptionFilter>,
    connected_at: std::time::SystemTime,
    last_activity: std::time::SystemTime,
}

/// Subscription filter for selective updates
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SubscriptionFilter {
    /// Filter by node types
    pub node_types: Option<Vec<String>>,
    /// Filter by languages
    pub languages: Option<Vec<String>>,
    /// Filter by file patterns
    pub file_patterns: Option<Vec<String>>,
    /// Filter by project paths
    pub project_paths: Option<Vec<String>>,
}

/// WebSocket message types
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type")]
pub enum WebSocketMessage {
    /// Client authentication
    Auth {
        token: Option<String>,
        user_id: Option<String>,
    },
    /// Subscribe to updates
    Subscribe { filters: Vec<SubscriptionFilter> },
    /// Unsubscribe from updates
    Unsubscribe { filter_ids: Vec<String> },
    /// Request current graph state
    GetGraphState { project_path: Option<String> },
    /// Request performance metrics
    GetMetrics,
    /// Ping for connection health
    Ping,
    /// Pong response
    Pong,
    /// Error message
    Error {
        code: String,
        message: String,
        details: Option<serde_json::Value>,
    },
}

/// Graph update message for broadcasting
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GraphUpdateMessage {
    /// Update ID
    pub id: Uuid,
    /// Timestamp
    pub timestamp: std::time::SystemTime,
    /// Type of update
    pub update_type: GraphUpdateType,
    /// Affected project path
    pub project_path: String,
    /// Update details
    pub details: GraphUpdateDetails,
    /// Performance metrics
    pub metrics: Option<UpdateMetrics>,
}

/// Type of graph update
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum GraphUpdateType {
    NodeAdded,
    NodeUpdated,
    NodeRemoved,
    EdgeAdded,
    EdgeUpdated,
    EdgeRemoved,
    FileChanged,
    ProjectRebuilt,
    RelationshipsUpdated,
}

/// Details of a graph update
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GraphUpdateDetails {
    /// Affected nodes
    pub nodes: Vec<NodeId>,
    /// Affected edges
    pub edges: Vec<EdgeId>,
    /// Changed files
    pub files: Vec<String>,
    /// Update summary
    pub summary: String,
    /// Additional metadata
    pub metadata: HashMap<String, serde_json::Value>,
}

/// Performance metrics for an update
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct UpdateMetrics {
    pub duration_ms: u64,
    pub nodes_processed: usize,
    pub edges_processed: usize,
    pub memory_usage_mb: f64,
}

/// WebSocket response messages
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type")]
pub enum WebSocketResponse {
    /// Authentication result
    AuthResult {
        success: bool,
        connection_id: Uuid,
        message: String,
    },
    /// Subscription confirmation
    SubscriptionResult {
        success: bool,
        subscription_id: String,
        message: String,
    },
    /// Graph state response
    GraphState {
        nodes: Vec<serde_json::Value>,
        edges: Vec<serde_json::Value>,
        metadata: HashMap<String, serde_json::Value>,
    },
    /// Performance metrics response
    Metrics {
        metrics: serde_json::Value,
        generated_at: std::time::SystemTime,
    },
    /// Real-time update
    Update(GraphUpdateMessage),
    /// Pong response
    Pong,
    /// Error response
    Error {
        code: String,
        message: String,
        details: Option<serde_json::Value>,
    },
}

impl WebSocketManager {
    /// Create a new WebSocket manager
    pub fn new(performance_monitor: Arc<PerformanceMonitor>) -> Self {
        let (update_sender, _) = broadcast::channel(1000);

        Self {
            connections: Arc::new(RwLock::new(HashMap::new())),
            update_sender,
            performance_monitor,
        }
    }

    /// Handle WebSocket upgrade
    pub async fn handle_websocket_upgrade(&self, ws: WebSocketUpgrade) -> Response {
        let connections = self.connections.clone();
        let update_sender = self.update_sender.clone();
        let performance_monitor = self.performance_monitor.clone();

        ws.on_upgrade(move |socket| {
            Self::handle_websocket_connection(
                socket,
                connections,
                update_sender,
                performance_monitor,
            )
        })
    }

    /// Handle individual WebSocket connection
    async fn handle_websocket_connection(
        socket: WebSocket,
        connections: Arc<RwLock<HashMap<Uuid, ConnectionInfo>>>,
        update_sender: broadcast::Sender<GraphUpdateMessage>,
        performance_monitor: Arc<PerformanceMonitor>,
    ) {
        let connection_id = Uuid::new_v4();
        let mut update_receiver = update_sender.subscribe();

        info!("New WebSocket connection: {}", connection_id);

        // Add connection to registry
        {
            let mut conns = connections.write().await;
            conns.insert(
                connection_id,
                ConnectionInfo {
                    id: connection_id,
                    user_id: None,
                    subscriptions: Vec::new(),
                    connected_at: std::time::SystemTime::now(),
                    last_activity: std::time::SystemTime::now(),
                },
            );
        }

        let (mut sender, mut receiver) = socket.split();

        // Create a channel for sending messages from incoming task to outgoing task
        let (msg_sender, mut msg_receiver) = tokio::sync::mpsc::unbounded_channel::<Message>();

        // Spawn task to handle incoming messages
        let connections_clone = connections.clone();
        let performance_monitor_clone = performance_monitor.clone();
        let incoming_task = tokio::spawn(async move {
            while let Some(msg) = receiver.next().await {
                match msg {
                    Ok(Message::Text(text)) => {
                        if let Err(e) = Self::handle_incoming_message(
                            &text,
                            connection_id,
                            &connections_clone,
                            &performance_monitor_clone,
                            &msg_sender,
                        )
                        .await
                        {
                            error!("Error handling WebSocket message: {}", e);
                        }
                    },
                    Ok(Message::Binary(_)) => {
                        warn!("Received binary message, ignoring");
                    },
                    Ok(Message::Close(_)) => {
                        info!("WebSocket connection closed: {}", connection_id);
                        break;
                    },
                    Err(e) => {
                        error!("WebSocket error: {}", e);
                        break;
                    },
                    _ => {},
                }
            }
        });

        // Spawn task to handle outgoing updates and direct messages
        let connections_for_outgoing = connections.clone();
        let outgoing_task = tokio::spawn(async move {
            loop {
                tokio::select! {
                    // Handle broadcast updates
                    update_result = update_receiver.recv() => {
                        match update_result {
                            Ok(update) => {
                                // Check if this connection should receive this update
                                let should_send = {
                                    let conns = connections_for_outgoing.read().await;
                                    if let Some(conn_info) = conns.get(&connection_id) {
                                        Self::should_send_update(&update, &conn_info.subscriptions)
                                    } else {
                                        false
                                    }
                                };

                                if should_send {
                                    let response = WebSocketResponse::Update(update);
                                    if let Ok(json) = serde_json::to_string(&response) {
                                        if sender.send(Message::Text(json)).await.is_err() {
                                            break;
                                        }
                                    }
                                }
                            }
                            Err(_) => break,
                        }
                    }
                    // Handle direct messages from incoming task
                    msg_result = msg_receiver.recv() => {
                        match msg_result {
                            Some(msg) => {
                                if sender.send(msg).await.is_err() {
                                    break;
                                }
                            }
                            None => break,
                        }
                    }
                }
            }
        });

        // Wait for either task to complete
        tokio::select! {
            _ = incoming_task => {},
            _ = outgoing_task => {},
        }

        // Clean up connection
        {
            let mut conns = connections.write().await;
            conns.remove(&connection_id);
        }

        info!("WebSocket connection cleaned up: {}", connection_id);
    }

    /// Handle incoming WebSocket message
    async fn handle_incoming_message(
        text: &str,
        connection_id: Uuid,
        connections: &Arc<RwLock<HashMap<Uuid, ConnectionInfo>>>,
        performance_monitor: &Arc<PerformanceMonitor>,
        sender: &tokio::sync::mpsc::UnboundedSender<Message>,
    ) -> Result<(), Box<dyn std::error::Error>> {
        let message: WebSocketMessage = serde_json::from_str(text)?;

        debug!("Received WebSocket message: {:?}", message);

        match message {
            WebSocketMessage::Auth { token: _, user_id } => {
                // Simple authentication (in production, validate token properly)
                let auth_success = user_id.is_some();

                // Update connection info with user ID
                {
                    let mut conns = connections.write().await;
                    if let Some(conn_info) = conns.get_mut(&connection_id) {
                        conn_info.user_id = user_id.clone();
                        conn_info.last_activity = std::time::SystemTime::now();
                    }
                }

                let response = WebSocketResponse::AuthResult {
                    success: auth_success,
                    connection_id,
                    message: if auth_success {
                        "Authentication successful".to_string()
                    } else {
                        "Authentication failed".to_string()
                    },
                };

                Self::send_response(sender, response).await?;
                info!(
                    "WebSocket authentication for connection: {} - Success: {}",
                    connection_id, auth_success
                );
            },

            WebSocketMessage::Subscribe { filters } => {
                // Update subscription filters
                {
                    let mut conns = connections.write().await;
                    if let Some(conn_info) = conns.get_mut(&connection_id) {
                        conn_info.subscriptions = filters;
                        conn_info.last_activity = std::time::SystemTime::now();
                    }
                }

                info!("Updated subscriptions for connection: {}", connection_id);
            },

            WebSocketMessage::GetMetrics => {
                // Send current performance metrics
                let report = performance_monitor.get_report().await;
                let _metrics_json = serde_json::to_value(&report)?;
                // TODO: Send metrics response back to client
                debug!("Sending metrics to connection: {}", connection_id);
            },

            WebSocketMessage::Ping => {
                // Update last activity
                {
                    let mut conns = connections.write().await;
                    if let Some(conn_info) = conns.get_mut(&connection_id) {
                        conn_info.last_activity = std::time::SystemTime::now();
                    }
                }

                // TODO: Send pong response
                debug!("Ping received from connection: {}", connection_id);
            },

            _ => {
                debug!("Unhandled WebSocket message type");
            },
        }

        Ok(())
    }

    /// Send a response message to the WebSocket client
    async fn send_response(
        sender: &tokio::sync::mpsc::UnboundedSender<Message>,
        response: WebSocketResponse,
    ) -> Result<(), Box<dyn std::error::Error>> {
        let json = serde_json::to_string(&response)?;
        sender.send(Message::Text(json))?;
        Ok(())
    }

    /// Check if an update should be sent to a connection based on subscriptions
    fn should_send_update(
        update: &GraphUpdateMessage,
        subscriptions: &[SubscriptionFilter],
    ) -> bool {
        if subscriptions.is_empty() {
            return true; // Send all updates if no filters
        }

        for filter in subscriptions {
            // Check project path filter
            if let Some(ref paths) = filter.project_paths {
                if !paths.iter().any(|path| update.project_path.contains(path)) {
                    continue;
                }
            }

            // Check file pattern filter
            if let Some(ref patterns) = filter.file_patterns {
                let matches_pattern = update
                    .details
                    .files
                    .iter()
                    .any(|file| patterns.iter().any(|pattern| file.contains(pattern)));
                if !matches_pattern {
                    continue;
                }
            }

            // If we get here, at least one filter matches
            return true;
        }

        false
    }

    /// Broadcast a graph update to all connected clients
    pub async fn broadcast_update(
        &self,
        update: GraphUpdateMessage,
    ) -> Result<(), Box<dyn std::error::Error>> {
        debug!("Broadcasting graph update: {:?}", update.id);

        if let Err(e) = self.update_sender.send(update) {
            warn!("Failed to broadcast update: {}", e);
        }

        Ok(())
    }

    /// Get connection statistics
    pub async fn get_connection_stats(&self) -> ConnectionStats {
        let connections = self.connections.read().await;

        ConnectionStats {
            total_connections: connections.len(),
            authenticated_connections: connections
                .values()
                .filter(|conn| conn.user_id.is_some())
                .count(),
            active_subscriptions: connections
                .values()
                .map(|conn| conn.subscriptions.len())
                .sum(),
        }
    }

    /// Broadcast node addition
    pub async fn broadcast_node_added(
        &self,
        project_path: String,
        node_ids: Vec<NodeId>,
        files: Vec<String>,
    ) -> Result<(), Box<dyn std::error::Error>> {
        let update = GraphUpdateMessage {
            id: Uuid::new_v4(),
            timestamp: std::time::SystemTime::now(),
            update_type: GraphUpdateType::NodeAdded,
            project_path,
            details: GraphUpdateDetails {
                nodes: node_ids,
                edges: Vec::new(),
                files,
                summary: "Nodes added to graph".to_string(),
                metadata: HashMap::new(),
            },
            metrics: None,
        };

        self.broadcast_update(update).await
    }

    /// Broadcast file change
    pub async fn broadcast_file_changed(
        &self,
        project_path: String,
        changed_files: Vec<String>,
        affected_nodes: Vec<NodeId>,
        affected_edges: Vec<EdgeId>,
    ) -> Result<(), Box<dyn std::error::Error>> {
        let update = GraphUpdateMessage {
            id: Uuid::new_v4(),
            timestamp: std::time::SystemTime::now(),
            update_type: GraphUpdateType::FileChanged,
            project_path,
            details: GraphUpdateDetails {
                nodes: affected_nodes,
                edges: affected_edges,
                files: changed_files,
                summary: "Files changed, graph updated".to_string(),
                metadata: HashMap::new(),
            },
            metrics: None,
        };

        self.broadcast_update(update).await
    }

    /// Broadcast project rebuild
    pub async fn broadcast_project_rebuilt(
        &self,
        project_path: String,
        total_nodes: usize,
        total_edges: usize,
        duration_ms: u64,
    ) -> Result<(), Box<dyn std::error::Error>> {
        let mut metadata = HashMap::new();
        metadata.insert(
            "total_nodes".to_string(),
            serde_json::Value::Number(total_nodes.into()),
        );
        metadata.insert(
            "total_edges".to_string(),
            serde_json::Value::Number(total_edges.into()),
        );
        metadata.insert(
            "duration_ms".to_string(),
            serde_json::Value::Number(duration_ms.into()),
        );

        let update = GraphUpdateMessage {
            id: Uuid::new_v4(),
            timestamp: std::time::SystemTime::now(),
            update_type: GraphUpdateType::ProjectRebuilt,
            project_path,
            details: GraphUpdateDetails {
                nodes: Vec::new(),
                edges: Vec::new(),
                files: Vec::new(),
                summary: format!(
                    "Project rebuilt: {} nodes, {} edges",
                    total_nodes, total_edges
                ),
                metadata,
            },
            metrics: Some(UpdateMetrics {
                duration_ms,
                nodes_processed: total_nodes,
                edges_processed: total_edges,
                memory_usage_mb: 0.0, // TODO: Get actual memory usage
            }),
        };

        self.broadcast_update(update).await
    }

    /// Clean up inactive connections
    pub async fn cleanup_inactive_connections(&self, timeout_seconds: u64) {
        let mut connections = self.connections.write().await;
        let now = std::time::SystemTime::now();
        let timeout_duration = std::time::Duration::from_secs(timeout_seconds);

        let inactive_connections: Vec<Uuid> = connections
            .iter()
            .filter_map(|(id, conn)| {
                if now.duration_since(conn.last_activity).unwrap_or_default() > timeout_duration {
                    Some(*id)
                } else {
                    None
                }
            })
            .collect();

        for conn_id in inactive_connections {
            connections.remove(&conn_id);
            info!("Cleaned up inactive WebSocket connection: {}", conn_id);
        }
    }
}

/// WebSocket connection statistics
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ConnectionStats {
    pub total_connections: usize,
    pub authenticated_connections: usize,
    pub active_subscriptions: usize,
}

use futures_util::{SinkExt, StreamExt};
